import numpy as np
import os
import pandas as pd
from sklearn.metrics import recall_score, precision_score, roc_curve, confusion_matrix
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split


class LogisticRegr(object):

    def __init__(self, N):
        self.w = np.random.randn(N, 1)
        self.b = 0.

    def predict(self, x):
        h = np.dot(x, self.w) + self.b
        h = self.sigmoid(h)
        return h

    @staticmethod
    def sigmoid(z):
        return 1. / (1. + np.e ** (- z))

    def cost(self, a, y):
        m = len(a)
        cost = (np.dot(y.T, np.log(a)) + np.dot(1 - y.T, np.log(1 - a)))[0][0] / m
        return - cost

    def dz(self, a, y):
        dzv = a - y
        return dzv

    def dw(self, a, y, x):
        m = len(a)
        dzv = self.dz(a, y)
        dwv = np.dot(x.T, dzv) / m
        return dwv

    def db(self, a, y):
        m = len(a)
        dbv = np.dot(np.ones((1, m), dtype=np.float64), self.dz(a, y)) / m
        return dbv

    @staticmethod
    def accuracy(y, a):
        acc = np.float64(
            np.equal(
                a > 0.5,
                y > 0.5
            )
        ).mean()
        return acc

    def metrics(self, y, a):
        acc = self.accuracy(y, a)
        precision = precision_score(y > 0.5, a > 0.5)
        recall = recall_score(y > 0.5, a > 0.5)
        return acc, precision, recall

    def confusion_matrix(self, y, a):
        TN = np.int64((y < 0.5) & (a < 0.5)).sum()
        FP = np.int64((y < 0.5) & (a > 0.5)).sum()
        FN = np.int64((y > 0.5) & (a < 0.5)).sum()
        TP = np.int64((y > 0.5) & (a > 0.5)).sum()
        mat = np.array([
            [TN, FP],
            [FN, TP],
        ])
        print(mat)
        print(confusion_matrix(y > 0.5, a > 0.5))

    def train(self, x, y, ITERS, ALPHA):
        j_his, acc_his, precision_his, recall_his = [], [], [], []
        group = int(np.ceil(ITERS / 20))
        for i in range(ITERS):
            a = self.predict(x)
            cost = self.cost(a, y)
            j_his.append(cost)
            if i == ITERS - 1 or i % group == 0:
                acc, precision, recall = self.metrics(y, a)
                print(f'#{i + 1}: cost = {cost}, acc = {acc}, precision = {precision}, recall = {recall}')
                acc_his.append(acc)
                precision_his.append(precision)
                recall_his.append(recall)
            dwv = self.dw(a, y, x)
            dbv = self.db(a, y)
            self.w -= ALPHA * dwv
            self.b -= ALPHA * dbv
        return np.float64(j_his), np.float64(acc_his), np.float64(precision_his), np.float64(recall_his), a

    def test(self, x, y):
        print('Testing ...')
        a = self.predict(x)
        cost = self.cost(a, y)
        acc, precision, recall = self.metrics(y, a)
        print(f'Test: cost = {cost}, acc = {acc}, precision = {precision}, recall = {recall}')
        return a, acc, precision, recall

    def precision_over_recall(self, y, a, split=100):
        y = y.ravel()
        a = a.ravel()
        idx = a.argsort()[::-1]
        y_sorted = y[idx]
        a_sorted = a[idx]
        len_y = len(y)
        len_positive = np.int64(y > 0.5).sum()
        len_unit = int(np.ceil(len_y / split))
        precisions, recalls = [], []
        for i in range(split):
            this_y = y_sorted[:(i + 1) * len_unit]
            this_a = a_sorted[:(i + 1) * len_unit]
            len_this = len(this_y)
            if len_this <= 0:
                break
            len_positive_this = np.int64(this_y > 0.5).sum()
            p = len_positive_this / len_this
            r = len_positive_this / len_positive
            precisions.append(p)
            recalls.append(r)
        return recalls, precisions


if '__main__' == __name__:

    np.random.seed(1)

    VER = 'v1.0'
    ALPHA = 1e-2
    ITERS = 2000
    BASE_DIR, FILE_NAME = os.path.split(__file__)
    path = '../../../large_data/ML1/hand_writing/imgX.txt'
    X_PATH = os.path.join(BASE_DIR, path)
    Y_PATH = os.path.join(BASE_DIR, path, '..', 'labely.txt')
    SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
    LOG_DIR = os.path.join(BASE_DIR, '_log', FILE_NAME, VER)

    x = np.loadtxt(X_PATH, delimiter=',').astype(np.float64)
    y = np.loadtxt(Y_PATH).astype(np.int64).reshape(-1, 1)
    y = np.int64(y == 5)
    print('x', x.shape, x.dtype)
    print('y', y.shape, y.dtype)

    y_pd = pd.Series(y.reshape(-1))
    print('y value counts:')
    print(y_pd.value_counts())

    x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.7, random_state=1, shuffle=True)

    logisticRegr = LogisticRegr(x.shape[1])
    j_his, acc_his, precision_his, recall_his, a_train = logisticRegr.train(x_train, y_train, ITERS, ALPHA)
    a_test, _, _, _ = logisticRegr.test(x_test, y_test)
    fpr, tpr, thresh = roc_curve(y_test, a_test)
    logisticRegr.confusion_matrix(y_test, a_test)
    recalls, precisions = logisticRegr.precision_over_recall(y_test, a_test)

    spr = 2
    spc = 3
    spn = 0
    plt.figure(figsize=(12, 8))
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('Cost')
    plt.plot(j_his)
    plt.grid()
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('Ratio')
    plt.plot(acc_his, label='accuracy')
    plt.plot(precision_his, label='precision')
    plt.plot(recall_his, label='recall')
    plt.grid()
    plt.legend()
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('ROC')
    plt.plot(fpr, tpr)
    plt.grid()
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('Precision on Recall')
    plt.plot(recalls, precisions)
    plt.grid()
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('Precision and Recall')
    plt.plot(precisions, label='precisions')
    plt.plot(recalls, label='recalls')
    plt.grid()
    plt.legend()
    plt.show()
